use thiserror::Error;

use stencil_tysys::{ArgRef, FieldRef, FnContext, LocalRef, ValueType};
use stencil_util::FunctionName;

use super::{
    function::MirFunction,
    var_tracker::{LocalVarTracker, VarError},
    ContainerData, CtorAssignments, FnKind, MirBlock, MirSignature, MirStatement, TaggedExpr,
    TypedExpr,
};

use crate::{typecheck_expr, AbortReason, AssignTarget, TypeError};

#[derive(Debug, Error)]
pub enum BuilderError {
    #[error("tried to insert statement after return")]
    StatementAfterReturn,

    #[error("inserted return type ({0:?}) did not match signature ({1:?})")]
    InvalidReturnValue(Option<ValueType>, Option<ValueType>),

    #[error("tried to insert conditional with invalid type")]
    CondInvalidType,

    #[error("tried to insert constructor in non-instance fn")]
    CtorInNonInstance,

    #[error("unknown field {0:?}")]
    UnknownField(FieldRef),

    #[error("tried to assign value {0:?} to ctor field of type {1:?}")]
    CtorTypeInvalid(ValueType, ValueType),

    #[error("var: {0}")]
    Var(#[from] VarError),

    #[error("typecheck: {0}")]
    Type(#[from] TypeError),
}

/// Builder for MIR functions.
struct FunctionBuilder {
    name: FunctionName,
    sig: MirSignature,
    kind: FnKind,

    /// The container this function might be an instance method of.  Provides
    /// necessary information for reasoning about fields.
    cont: Option<ContainerData>,

    /// Variable tracker used to provide env types to exprs and track var refs.
    var_tracker: LocalVarTracker,
}

impl FunctionBuilder {
    fn new(
        name: FunctionName,
        sig: MirSignature,
        kind: FnKind,
        cont: Option<ContainerData>,
    ) -> Self {
        let var_tracker = LocalVarTracker::new();
        Self {
            name,
            sig,
            kind,
            cont,
            var_tracker,
        }
    }

    /// Returns the builder context used to generate blocks.
    pub fn get_fn_context(&mut self) -> BuilderContext<'_> {
        BuilderContext {
            sig: &self.sig,
            cont: self.cont.as_ref(),
            var_tracker: &mut self.var_tracker,
        }
    }

    /// Completes the function with a root code block.
    pub fn finish(self, root_block: MirBlock) -> MirFunction {
        MirFunction::new(self.name, self.sig, self.kind, root_block)
    }
}

/// Context necessary for generating blocks within a function that respect env
/// types and track locals appropriately.
///
/// Also supports being used as a typechecker.
pub struct BuilderContext<'c> {
    sig: &'c MirSignature,
    cont: Option<&'c ContainerData>,
    var_tracker: &'c mut LocalVarTracker,
}

impl<'c> BuilderContext<'c> {
    /// Checks an expression, returning a typed copy of it.
    pub fn check_expr<T>(&self, expr: &TaggedExpr<T>) -> Result<TypedExpr, TypeError> {
        typecheck_expr(&expr.to_int_tag(), self)
    }
}

impl<'c> FnContext for BuilderContext<'c> {
    fn is_container(&self) -> bool {
        self.cont.is_some()
    }

    fn num_args(&self) -> u16 {
        self.sig.arg_tys().len() as u16
    }

    fn get_arg_type(&self, arg: ArgRef) -> Option<ValueType> {
        self.sig.arg_tys().get(arg.inner() as usize).cloned()
    }

    fn get_local_type(&self, local: LocalRef) -> Option<ValueType> {
        self.var_tracker.get_local_ty(local).cloned()
    }

    fn num_instance_fields(&self) -> u16 {
        self.cont
            .map(|c| c.fields().len() as u16)
            .unwrap_or_default()
    }

    fn get_instance_field_type(&self, field: FieldRef) -> Option<ValueType> {
        self.cont
            .and_then(|c| c.get_field(field))
            .map(|cf| cf.ty().clone())
    }
}

/// Builder for MIR blocks.
struct BlockBuilder<'b> {
    /// Backing context tracking types and locals.
    ctx: &'b mut BuilderContext<'b>,

    /// List of entries we're building.
    entries: Vec<MirStatement>,

    /// If the builder has had a return inserted in this block.
    did_return: bool,
}

impl<'b> BlockBuilder<'b> {
    pub fn new(ctx: &'b mut BuilderContext<'b>) -> Self {
        Self {
            ctx,
            entries: Vec::new(),
            did_return: false,
        }
    }

    /// Finishes the block, producing a [`MirBlock`].
    pub fn finish(self) -> MirBlock {
        MirBlock::new(self.entries)
    }

    /// Checks to see if we already returned.  If we did, then there's no sense
    /// in inserting most new statements and should call `.finish()`.
    fn check_unreturned(&self) -> Result<(), BuilderError> {
        if self.did_return {
            Err(BuilderError::StatementAfterReturn)
        } else {
            Ok(())
        }
    }

    fn set_returned(&mut self) {
        self.did_return = true;
    }

    /// Checks if a return passed return type matches the signature.
    fn check_return_type(&self, rty: Option<&ValueType>) -> Result<(), BuilderError> {
        let exp_rty = self.ctx.sig.ret_ty();
        if rty != exp_rty {
            return Err(BuilderError::InvalidReturnValue(
                rty.cloned(),
                exp_rty.cloned(),
            ));
        } else {
            Ok(())
        }
    }

    /// Typechecks an expression.
    pub fn check_expr<T>(&self, expr: &TaggedExpr<T>) -> Result<TypedExpr, TypeError> {
        // FIXME make this ignore locals declared in expressions above this one
        self.ctx.check_expr(expr)
    }

    /// Inserts a new statement.
    fn insert_statement(&mut self, stmt: MirStatement) {
        self.entries.push(stmt);
    }

    /// Puts a comment.
    pub fn put_comment(&mut self, comment: String) {
        self.insert_statement(MirStatement::Comment(comment));
    }

    /// Puts an assignment to a local directly, simply.
    fn put_raw_assign(&mut self, lr: LocalRef, expr: TypedExpr) -> Result<(), BuilderError> {
        // TODO avoid clone
        self.ctx
            .var_tracker
            .reassign_local(lr, expr.tag().clone())?;
        self.insert_statement(MirStatement::Assign(AssignTarget::Local(lr), expr));
        Ok(())
    }

    /// Assigns an local with a new expression value, returning a ref to the
    /// new local created.
    pub fn put_assign_new_local(&mut self, expr: TypedExpr) -> Result<LocalRef, BuilderError> {
        let lr = self.ctx.var_tracker.declare_local(expr.tag().clone());
        self.put_raw_assign(lr, expr)?;
        Ok(lr)
    }

    /// Puts a void return.
    pub fn put_void_return(&mut self) -> Result<(), BuilderError> {
        self.check_unreturned()?;
        self.check_return_type(None)?;
        self.insert_statement(MirStatement::ReturnVoid);
        self.set_returned();
        Ok(())
    }

    /// Puts a value return expression.
    pub fn put_value_return(&mut self, expr: TypedExpr) -> Result<(), BuilderError> {
        self.check_unreturned()?;
        self.check_return_type(Some(expr.tag()))?;
        self.insert_statement(MirStatement::ReturnValue(expr));
        self.set_returned();
        Ok(())
    }

    /// Puts a constructor expression, which is treated like a return.
    ///
    /// This *should* be at the root block, but but we don't have a great way to
    /// ensure that here.
    pub fn put_constructor(&mut self, ctor: CtorAssignments) -> Result<(), BuilderError> {
        self.check_unreturned()?;

        // Check that we are a container that can probably be constructed here.
        if self.ctx.is_container() {
            return Err(BuilderError::CtorInNonInstance);
        }

        // Check the assignments match field types.
        for assign in ctor.assignments() {
            let ety = assign.expr().tag();

            let fty = self
                .ctx
                .get_instance_field_type(assign.field())
                .ok_or(BuilderError::UnknownField(assign.field()))?;

            if &fty != ety {
                return Err(BuilderError::CtorTypeInvalid(ety.clone(), fty));
            }
        }

        self.insert_statement(MirStatement::Construct(ctor));
        self.set_returned();

        Ok(())
    }

    /// Puts an `If` block that just aborts with some reason.
    ///
    /// This is a convenience function that avoids doing the work for
    /// constructing a new block builder just for the abort.
    pub fn put_conditional_abort(
        &mut self,
        cond_expr: TypedExpr,
        reason: AbortReason,
    ) -> Result<(), BuilderError> {
        let block = MirBlock::new(vec![MirStatement::Abort(reason)]);
        let stmt = MirStatement::If(cond_expr, Box::new(block));
        self.insert_statement(stmt);
        Ok(())
    }

    /// Creates a "child" builder instance that maintains the local var context.
    ///
    /// This can be used to construct child blocks for things like branches,
    /// loops, etc.  This instance has to be dropped before we can use this
    /// context again.
    pub fn create_child_builder<'cb>(&'cb mut self) -> BlockBuilder<'cb>
    where
        'cb: 'b,
    {
        Self::new(self.ctx)
    }
}
